"""来自论文Crater-based attitude and position estimation for planetary exploration with weighted measurement uncertainty
  不确定度的计算使用第3页
  加权POSIT算法的计算使用第7页
"""

from .base import PoseEstimator, center_ellipse
import numpy as np


class CP(PoseEstimator):
    def __init__(self, K, use_pnp_initial=True, *args, **kwargs):
        super().__init__(K, *args, **kwargs)
        self.use_pnp_initial = use_pnp_initial

    def __name__(self):
        return "cp"

    def forward(self, C_3d, C_2d, *args, **ignore_kwargs):

        x_cnt = []
        X_cnt = []
        for c_3d, c_2d in zip(C_3d, C_2d):
            # 求解相机椭圆中心：
            x_cnt.append(center_ellipse(c_2d))
            X_cnt.append(center_ellipse(c_3d))
        # 求解相机位姿
        X_cnt = np.pad(X_cnt, ((0, 0), (0, 1)), "constant", constant_values=0)
        x_cnt = np.array(x_cnt)
        # 调用opencv的库实现加权POSIT算法
        return self.posit_algorithm(X_cnt, x_cnt)

    def posit_algorithm(
        self,
        points_3d,
        points_2d,
        max_iterations=1000,
        tol=1e-8,
    ):
        """POSIT算法求解位姿"""
        R = np.eye(3)
        t = np.zeros(2)
        # 初始猜测
        lamda = np.ones(len(points_2d))
        # 迭代求解
        for _ in range(max_iterations):
            # 构建线性方程组
            A = np.block(
                [
                    [
                        points_3d,
                        np.zeros_like(points_3d),
                        np.array([[1, 0]] * len(points_3d)),
                    ],
                    [
                        np.zeros_like(points_3d),
                        points_3d,
                        np.array([[0, 1]] * len(points_3d)),
                    ],
                ]
            )
            b = np.block([lamda * points_2d[:, 0], lamda * points_2d[:, 1]])
            x = np.linalg.lstsq(A, b, rcond=None)[0]
            # s * c1, s * c2, s * t = x
            # 计算更新的比例因子
            s_now = np.sqrt(np.linalg.norm(x[:3]) * np.linalg.norm(x[3:6]))
            if np.linalg.norm(x[6:] / s_now - t) < tol:
                break
            # 计算当前迭代步的参数
            t = x[6:] / s_now
            Z_C = self.K[0, 0] / s_now
            R[0] = x[:3] / s_now
            R[1] = x[3:6] / s_now
            R[2] = np.cross(R[0], R[1])
            # 计算当前迭代数的lambda
            lamda = points_3d @ R[2] / Z_C + 1
        # 计算误差
        error = np.sqrt(np.power(A @ x - b, 2).mean())
        return True, R, np.block([t, Z_C])
